import os
import glob
import random
import time
import csv
import cv2
import numpy as np
from PyQt5.QtCore import *
from PyQt5.QtWidgets import QWidget
from PyQt5.Qt import pyqtSignal
import win32file
import shutil

# from AI.AI_Segmention import AI_Segment
from AI_Densenet.Compoment_Classify import Compoment_Classify
from AI_RC.RC_classify import C_Classify,R_Classify,All_Classify,D_Classify,Q_Classify
from AI_OCR .AI_OCR import AI_OCR
from AI_ICCN.IC_CN_YJ_classify import YJ_Det
from AI_ICCN.ICCN_Det import ICCN_Det
from utils.ptt_reader import read_ptt
##test push
def is_used(file_name):
    try:
        vHandle = win32file.CreateFile(file_name, win32file.GENERIC_READ, 0, None, win32file.OPEN_EXISTING, win32file.FILE_ATTRIBUTE_NORMAL, None)
        return int(vHandle) == win32file.INVALID_HANDLE_VALUE
    except:
        return True
    finally:
        try:
            win32file.CloseHandle(vHandle)
        except:
            pass


class WorkThread(QThread):
    # 自定义信号对象。参数str就代表这个信号可以传一个字符串
    trigger = pyqtSignal(str)
    image_sig = pyqtSignal(list)
    aoi_image_sig = pyqtSignal(list)
    _3d_image_sig = pyqtSignal(list)
    clear_sig = pyqtSignal()
    aoi_row_sig = pyqtSignal(list)
    ai_row_sig = pyqtSignal(list)
    aoi_box_sig = pyqtSignal(list)
    def __int__(self):
        # 初始化函数
        super().__init__()

    def load_image(self,image_path,det_img_box_path='',det_img_contour_path ='',box = []):
        if self.product_type == 'PemtronSPI':
            if os.path.exists(image_path):
                print(image_path)
                image = cv2.imdecode(np.fromfile(image_path,dtype= np.uint8),-1)
                # time.sleep(1)
                # out_image = self.segment.start_seg(image)
                out_image = image
                self.image_sig.emit([image,out_image])
                if self.box_save:
                    # merge_img = cv2.merge([image,image,image])
                    cv2.rectangle(image,(box[0],box[1]),(box[2],box[3]),(2,2,255),2)
                    # cv2.imshow('a',merge_img)
                    # cv2.waitKey(0)tifs = tifs.unsqueeze(1)
                    cv2.imencode('.jpg', image)[1].tofile(det_img_box_path)
                # if self.contour_save:
                #     cv2.imencode('.jpg', out_image)[1].tofile(det_img_contour_path)
                # res = random.sample(['22','23','1'],1)[0]
                res = '22'
                return res
        elif self.product_type == 'PemtronAOI':
            if os.path.exists(image_path):
                print(image_path)
                image_AC_path = image_path.replace('.jpg','_AC.jpg')
                image_ptt_path = image_path.replace('.jpg','.ptt')

                image = cv2.imdecode(np.fromfile(image_path,dtype= np.uint8),-1)
                # shutil.copyfile(image_path,det_img_path)
                if os.path.exists(image_AC_path):
                    image_AC = cv2.imdecode(np.fromfile(image_AC_path,dtype= np.uint8),-1)
                    # det_img_AC_path = det_img_path.replace('.jpg','_AC.jpg')
                    # shutil.copyfile(image_AC_path,det_img_AC_path)
                else:
                    image_AC = image
                if os.path.exists(image_ptt_path):
                    # det_img_ptt_path = det_img_path.replace('.jpg', '.ptt')
                    # shutil.copyfile(image_ptt_path, det_img_ptt_path)
                    heatmap,texture = read_ptt(image_ptt_path)
                    self._3d_image_sig.emit([heatmap, texture])
                self.aoi_image_sig.emit([image,image_AC,image])
                res = random.sample(['AIOK', 'AING'], 1)[0]
                # res = 'AING'
                return res
                # out_image = self.segment.start_seg(image)

            pass

    def find_positions(self,data, pattern):
        positions = []
        position = data.find(pattern)
        while position != -1:
            positions.append(position)
            position = data.find(pattern, position + 1)
        return positions

    def get_byte_string(self,file_path):
        with open(file_path, 'rb') as f:
            data = f.read(100)
        pattern = b'\xFF\xFE\xFF'
        positions = self.find_positions(data, pattern)
        start = positions[1] + 3
        end = positions[2]
        find_code = data[start:end]
        find_string = find_code.decode()
        txt_nmae = [i for i in find_string if i.isalnum()]
        txt_nmae = ''.join(txt_nmae)
        txt_nmae.replace('I','1')
        print('MODEL:',txt_nmae)
        return txt_nmae

    def insert_data(self,AOI_Folder,AI_Folder,AIOK_Folder,csv_paths,product_type,ai_models,box_save,contour_save,resolutiion,is_save_img,is_change_result,is_save_aiok,OCR_Root,POT_Root):
        self.AOI_Folder = AOI_Folder
        self.AI_Folder = AI_Folder
        self.csv_paths = csv_paths
        self.OCR_Root = OCR_Root
        self.POT_Root = POT_Root
        self.AIOK_Folder = AIOK_Folder
        self.product_type = product_type
        self.segment = None
        self.box_save = box_save
        self.contour_save = contour_save
        self.resolutiion = resolutiion
        self.is_save_img = is_save_img
        self.is_change_result = is_change_result
        self.is_save_aiok = is_save_aiok
        self.com_classity = ai_models[0]
        self.c_classify = ai_models[1]
        self.r_classify = ai_models[2]
        self.iccn_det = ai_models[3]
        self.d_classify = ai_models[4]
        self.q_classify = ai_models[5]
        self.ai_ocr = ai_models[6]
        self.yj_det  = ai_models[7]

    def get_insp_array_dict(self,insp_array_name):
        """
            获取 insp_array 的数据
        """
        with open('./config/TempInspResult_Root.txt','r',encoding='utf-8') as f:
            row = f.readlines()
        TempInspResult_Root = row[0].strip('\n')
        insp_array_path = os.path.join(TempInspResult_Root,insp_array_name)
        if not os.path.exists(insp_array_path):
            print('{} is not exists'.format(insp_array_path))
            return
        insp_array_dict = {}
        fp_insp_array = open(insp_array_path,"r",encoding="utf-8")
        line_2 = fp_insp_array.readline()
        while line_2:
            line_2 = fp_insp_array.readline()
            if line_2 != "":
                insp_array_dict[line_2.split(",")[1]] = line_2.split(",")[2]
        print(insp_array_dict)
        return insp_array_dict
    def run(self):
        i=0
        # start_time = time.localtime()
        #重写线程执行的run函数
        if self.product_type == 'PemtronSPI':
            csv_path = self.csv_paths[0]
            timefolder, base_name = os.path.split(csv_path)
            ai_folder = timefolder+'/AI'
            if not os.path.exists(ai_folder):
                os.makedirs(ai_folder)

            dst_csv_path = os.path.join(ai_folder, base_name)
        else:
            start_time = time.time()

            csv_path = self.csv_paths[0]
            timefolder,base_name = os.path.split(csv_path)
            timefolder1 = os.path.split(timefolder)[0]
            dst_csv_folder  = timefolder1.replace(self.AOI_Folder,self.AI_Folder)
            if not os.path.exists(dst_csv_folder):
                os.makedirs(dst_csv_folder)
            dst_csv_path = os.path.join(dst_csv_folder,base_name)

        while True:
            if is_used(dst_csv_path):
                print(dst_csv_path, 'not copy done!')
                time.sleep(0.5)
            else:
                break
        data = []
        # self.csv_pth_single.emit(csv_path)
        image_rows = {}
        ocr_rows = {}
        refid_rows = {}
        with open(dst_csv_path,encoding='utf-8') as csvfile:
            csv_reader = csv.reader(csvfile)  # 使用csv.reader读取csvfile中的文件
            for idx,row in enumerate(csv_reader):
                if idx == 0:
                    data.append(row)
                    continue
                date,groupname,slavename,arr_id, partcode,pad_no, aoi_res,ng_name ,angle =row[0],row[1],row[3], row[4],row[6], row[7], row[15],row[10],row[17]
                insp_array_path = os.path.join(groupname,slavename,date,"insp_array.txt")
                try:
                    REFID = row[16]
                    REFID = REFID.rstrip('0123456789-')
                except Exception as e:
                    REFID = ''
                arr_key = '{}@{}'.format(int(arr_id),int(pad_no))

                image_path = os.path.join(timefolder,"NGPartImage","{}.jpg".format(arr_key))
                image = cv2.imdecode(np.fromfile(image_path,dtype= np.uint8),-1)
                image_rows[arr_key] = image
                refid_rows[arr_key] = REFID
                data.append(row)
                if REFID in ['D','U','R','Q'] or ng_name == 'OCR':
                    # pot_path =
                    ocr_rows[arr_key] = [image,date,slavename,angle,REFID]
            ocr_res = self.ai_ocr.start_ocr(ocr_rows,self.POT_Root)
            output = self.com_classity.start_classify(image_rows)
            list_R=[]
            list_C=[]
            list_IC=[]
            list_CN=[]
            list_D=[]
            list_Q=[]
            list_NG = []
            ng_res = {}
            list_OT = []
            for key,value in output.items():
                # print(key)
                if value=='Konghanpan':
                    list_NG.append([key, value, image_rows[key]])
                    ng_res[key] = 'NG'
                else:
                    if refid_rows[key]=='':
                        if value == "R":
                            list_R.append([key,value,image_rows[key]])
                        elif value == "C":
                            list_C.append([key,value,image_rows[key]])
                        elif value == "IC":
                            list_IC.append([key,value,image_rows[key]])
                        elif value == "CN":
                            list_CN.append([key,value,image_rows[key]])
                        elif value == 'D':
                            list_D.append([key,value,image_rows[key]])
                        elif value == 'Q':
                            list_Q.append([key, value, image_rows[key]])
                        elif value == 'Other_NG':
                            list_NG.append([key,value,image_rows[key]])
                            ng_res[key] = 'NG'
                        else:
                            list_OT.append([key,value,image_rows[key]])
                            ng_res[key] = 'NG'
                    elif refid_rows[key] == 'R' and value=='R':
                        list_R.append([key,value,image_rows[key]])
                    elif refid_rows[key] == "C" and value=='C':
                        list_C.append([key,value,image_rows[key]])
                    elif refid_rows[key] == "U" and value=='IC':
                        list_IC.append([key,value,image_rows[key]])
                    elif refid_rows[key] in["X","J"]  and value=='CN':
                        list_IC.append([key,value,image_rows[key]])
                    elif refid_rows[key] == 'D':
                        if value == 'D':
                            list_D.append([key,value,image_rows[key]])
                        elif value == 'Q':
                            list_Q.append([key,value,image_rows[key]])
                    elif refid_rows[key] == 'Q' and value=='Q':
                        list_Q.append([key, value, image_rows[key]])
                    else:
                        list_OT.append([key, value, image_rows[key]])
                        ng_res[key] = 'NG'

            c_res = self.c_classify.start_classify(list_C,ng_res)
            r_res = self.r_classify.start_classify(list_R,ng_res)
            d_res = self.d_classify.start_classify(list_D,ng_res)
            q_res = self.q_classify.start_classify(list_Q,ng_res)
            ic_res = self.yj_det.start_classify(list_IC,ng_res,os.path.join(timefolder,"NGPartImage"))
            ic_res_2 = self.iccn_det.start_det(list_IC,ng_res)


            f = open(dst_csv_path, 'w', newline="", encoding='utf-8')
            # if os.path.exists(insp_array_path):
            insp_array = self.get_insp_array_dict(insp_array_path)
            # else:
            #     insp_array = None
            csv_writer = csv.writer(f)
            for i,row in enumerate(data):
                date,groupname,arr_id,partcode, pad_no, aoi_res,ng_name,refid  = row[0],row[1],row[4], row[6],row[7], row[15],row[10],row[16]

                if i ==0:
                    csv_writer.writerow(row)
                    continue
                if insp_array is not None:
                    new_img_name = "{}_{}_{}_{}_{}.jpg".format(date, insp_array[arr_id], key, refid, ng_name)
                else:
                    new_img_name = "{}_{}_{}_{}_{}.jpg".format(date, partcode, key, refid, ng_name)

                arr_key = '{}@{}'.format(arr_id, pad_no)
                if ng_name == 'OCR':
                    row[15] = 'AING'
                    #比较字符串
                    job_name, partcode = row[2], row[6]
                    psr_paths = glob.glob(os.path.join(self.OCR_Root,groupname,job_name,'Master/POCRFont/{}*.psr'.format(partcode)))
                    for psr_path in psr_paths:
                        print(psr_path)
                        muban_txt = self.get_byte_string(psr_path)
                        ocr_txts = ocr_res[arr_key]
                        print("OCR Result:",ocr_txts)
                        for ocr_txt in ocr_txts:
                            if muban_txt in ocr_txt:
                                row[15] = 'AIOK'
                                if self.is_save_aiok:


                                    aoijpg_path = os.path.join(self.AOI_Folder,date,"NGPartImage","{}.jpg".format(arr_key))
                                    aoiptt_path = aoijpg_path.replace('jpg','ptt')
                                    aoi_ac_path = aoijpg_path.replace('.jpg','_AC.jpg')
                                    aiokjpg_savepath = os.path.join(self.AIOK_Folder,new_img_name)
                                    aiokptt_savepath = aiokjpg_savepath.replace('jpg','ptt')
                                    aiok_ac_save_path = aiokjpg_savepath.replace('.jpg','_AC.jpg')
                                    try:
                                        shutil.copyfile(aoijpg_path,aiokjpg_savepath)
                                        shutil.copyfile(aoiptt_path, aiokptt_savepath)
                                        shutil.copyfile(aoi_ac_path, aiok_ac_save_path)
                                    except Exception as e:
                                        print(e)
                                continue


                    # row[15] = 'AING'
                    csv_writer.writerow(row)
                    continue
                # if ng_name in ['M_TILT','L_LIFT']:
                #     row[15] = 'AING'
                #     csv_writer.writerow(row)
                #     continue
                if ng_res.get(arr_key) is not None:
                    # print(arr_key)
                    row[15] = 'AING'
                    csv_writer.writerow(row)
                else:

                    # 比较字符串
                    groupname, job_name, partcode = row[1], row[2], row[6]
                    psr_paths = glob.glob(
                        os.path.join(self.OCR_Root, groupname, job_name, 'Master/POCRFont/{}*.psr'.format(partcode)))
                    if len(psr_paths):
                        print("AI OCR Det****************")
                        for psr_path in psr_paths:
                            if not os.path.exists(psr_path):
                                row[15] = 'AIOK'
                                if self.is_save_aiok:
                                    aoijpg_path = os.path.join(self.AOI_Folder,date,"NGPartImage","{}.jpg".format(arr_key))
                                    aoiptt_path = aoijpg_path.replace('jpg','ptt')
                                    aoi_ac_path = aoijpg_path.replace('.jpg','_AC.jpg')
                                    aiokjpg_savepath = os.path.join(self.AIOK_Folder,new_img_name)
                                    aiokptt_savepath = aiokjpg_savepath.replace('jpg','ptt')
                                    aiok_ac_save_path = aiokjpg_savepath.replace('.jpg','_AC.jpg')
                                    try:
                                        shutil.copyfile(aoijpg_path,aiokjpg_savepath)
                                        shutil.copyfile(aoiptt_path, aiokptt_savepath)
                                        shutil.copyfile(aoi_ac_path, aiok_ac_save_path)
                                    except Exception as e:
                                        print(e)
                            else:
                                row[15] = 'AING'
                                print(psr_path)
                                muban_txt = self.get_byte_string(psr_path)
                                ocr_txts = ocr_res[arr_key]
                                print("OCR Result:", ocr_txts)
                                for ocr_txt in ocr_txts:
                                    if muban_txt in ocr_txt:
                                        row[15] = 'AIOK'
                                        if self.is_save_aiok:
                                            aoijpg_path = os.path.join(self.AOI_Folder, date, "NGPartImage",
                                                                       "{}.jpg".format(arr_key))
                                            aoiptt_path = aoijpg_path.replace('jpg', 'ptt')
                                            aoi_ac_path = aoijpg_path.replace('.jpg', '_AC.jpg')
                                            aiokjpg_savepath = os.path.join(self.AIOK_Folder, new_img_name)
                                            aiokptt_savepath = aiokjpg_savepath.replace('jpg', 'ptt')
                                            aiok_ac_save_path = aiokjpg_savepath.replace('.jpg', '_AC.jpg')
                                            try:
                                                shutil.copyfile(aoijpg_path, aiokjpg_savepath)
                                                shutil.copyfile(aoiptt_path, aiokptt_savepath)
                                                shutil.copyfile(aoi_ac_path, aiok_ac_save_path)
                                            except Exception as e:
                                                print(e)
                                        continue
                    else:
                        row[15] = 'AIOK'
                        if self.is_save_aiok:
                            aoijpg_path = os.path.join(self.AOI_Folder, date, "NGPartImage", "{}.jpg".format(arr_key))
                            aoiptt_path = aoijpg_path.replace('jpg', 'ptt')
                            aoi_ac_path = aoijpg_path.replace('.jpg', '_AC.jpg')
                            aiokjpg_savepath = os.path.join(self.AIOK_Folder, new_img_name)
                            aiokptt_savepath = aiokjpg_savepath.replace('jpg', 'ptt')
                            aiok_ac_save_path = aiokjpg_savepath.replace('.jpg', '_AC.jpg')
                            try:
                                shutil.copyfile(aoijpg_path, aiokjpg_savepath)
                                shutil.copyfile(aoiptt_path, aiokptt_savepath)
                                shutil.copyfile(aoi_ac_path, aiok_ac_save_path)
                            except Exception as e:
                                print(e)
                    csv_writer.writerow(row)
            f.close()
            # 通过自定义信号把待显示的字符串传递给槽函数
            # self.trigger.emit(str(i))
            # end_time = time.localtime()
            # process_time = 1000*(end_time-start_time)
            end_time = time.time()
            print("t={}s".format(end_time - start_time))
            self.ai_row_sig.emit(['0', '0', 'Done!'])



class Win_FileWatcher(QWidget):
    trigger = pyqtSignal(str)
    clear_sig = pyqtSignal()
    coll__num_sig = pyqtSignal(int)
    aoi_row_sig = pyqtSignal(list)
    ai_row_sig = pyqtSignal(list)
    image_sig = pyqtSignal(list)
    aoi_image_sig = pyqtSignal(list)
    _3d_image_sig = pyqtSignal(list)
    aoi_box_sig = pyqtSignal(list)
    # addsin1 = pyqtSignal(str, str, str)
    def __init__(self):
        super(Win_FileWatcher,self).__init__()

        self.timer = QTimer()

        self.test_result = {"1_1":'22',
                       "2_8":'23',
                       "3_3":"23"}
        self.box_save = True
        self.contour_save = True
        self.resolution = 15.0
        self.is_collect = False
        self.is_change_res = True
        # self.AI_Segment = AI_Segment()
        self.AI_Segment = None
        self.collect_num = 0
        self.com_classity = All_Classify()
        self.c_classify = C_Classify()
        self.r_classify = R_Classify()
        self.d_classify = D_Classify()
        self.iccn_det = ICCN_Det()
        self.q_classify = Q_Classify()
        self.ai_ocr= AI_OCR()
        self.yj_det = YJ_Det()
    @classmethod
    def instance(cls):
        if not hasattr(Win_FileWatcher, "_instance"):
            Win_FileWatcher._instance = Win_FileWatcher()
        return Win_FileWatcher._instance

    def add_watchlist(self,*args): #,AOI_Folder,AI_Folder,Img_Connection_Folder,product_type,OCR_Root):
        self.AOI_Folder = args[0] #AOI_Folder
        self.AI_Folder = args[1]
        self.Img_Connection_Folder = args[2]
        self.product_type = args[3]
        self.OCR_Root = args[4]
        self.POT_Root = args[5]
        self.AIOK_Folder = args[6]
        self.is_save_aiok = args[7]
        if not os.path.exists(self.AI_Folder) and self.AOI_Folder !='':

            os.makedirs(self.AI_Folder)
        # self.trainImage_Path = os.path.join(self.AI_Data_Path, 'trainImage')

    def Collection_State(self,is_collected):
        self.is_collect = is_collected

    def Result_Change_State(self,is_changed):
        self.is_change_res = is_changed
    def Start(self):
        self.timerInterval = 0.1
        self.timer.timeout.connect(self.Update)
        self.timer.start(1000)
    def Update(self):
        if self.product_type=="PemtronSPI":
            csv_paths = glob.glob(os.path.join(self.AOI_Folder, '*/*.csv'))
        elif self.product_type == 'PemtronAOI':
            csv_paths = glob.glob(os.path.join(self.AOI_Folder, '*/*.csv'))
        if len(csv_paths):


            csv_path = csv_paths[0]
            if self.is_collect:
                src_folder = os.path.dirname(csv_path)
                collect_folder = src_folder.replace(self.AOI_Folder, self.Img_Connection_Folder)
                try:
                    shutil.copytree(src_folder, collect_folder)
                    self.collect_num +=1
                    self.collect_num.emit(self.collect_num)
                except Exception as e:
                    print(e)
            timefolder, base_name = os.path.split(csv_path)
            p,t = os.path.split(timefolder)
            copy_folder = os.path.join(p+'_AI',t)
            if not os.path.exists(copy_folder):
                os.makedirs(copy_folder)
            shutil.copyfile(csv_path, os.path.join(copy_folder,base_name))
            if self.product_type == 'PemtronSPI':


                ai_folder = timefolder + '/AI'
                if not os.path.exists(ai_folder):
                    os.makedirs(ai_folder)

                dst_csv_path = os.path.join(ai_folder, base_name)
            else:
                timefolder1 = os.path.split(timefolder)[0]
                # copy_folder = timefolder1.replace(self.AOI_Folder, self.AI_Folder)
                # if not os.path.exists(copy_folder):
                #     os.makedirs(copy_folder)
                # shutil.copyfile(csv_path, os.path.join(copy_folder, base_name))
                dst_csv_folder = timefolder1.replace(self.AOI_Folder, self.AI_Folder)
                if not os.path.exists(dst_csv_folder):
                    os.makedirs(dst_csv_folder)
                dst_csv_path = os.path.join(dst_csv_folder, base_name)
            shutil.move(csv_path, dst_csv_path)
            print("_____________________________")
            self.clear_sig.emit()
            self.work = WorkThread()
            self.work.insert_data(self.AOI_Folder,self.AI_Folder,self.AIOK_Folder,
                                  csv_paths,self.product_type,
                                  [self.com_classity,self.c_classify,self.r_classify,
                                   self.iccn_det,self.d_classify,self.q_classify,
                                   self.ai_ocr,self.yj_det],
                                   self.box_save,
                                  self.contour_save,self.resolution,self.is_collect,self.is_change_res,self.is_save_aiok,
                                  self.OCR_Root,self.POT_Root)
            self.work.start()
            # 线程自定义信号连接的槽函数
            self.work.trigger.connect(self.display)
            self.work.image_sig.connect(self.show_img)
            self.work.aoi_image_sig.connect(self.show_aoi_img)
            self.work.aoi_row_sig.connect(self.aoi_tablewidget)
            self.work.ai_row_sig.connect(self.ai_tablewidget)
            self.work.clear_sig.connect(self.clear_widget)
            self.work.aoi_box_sig.connect(self.roi_groupbox)
            self.work._3d_image_sig.connect(self.ThreeD_image)



    def display(self, str):
        self.trigger.emit(str)
        self.work.quit()
        self.work.wait()
        del self.work

    def set_save_state(self,box_save,contour_save,is_collect,is_change_res,is_aiok):
        self.box_save = box_save
        self.contour_save = contour_save
        self.is_collect = is_collect
        self.is_change_res = is_change_res
        self.is_save_aiok = is_aiok
        # self.resolution = resolution

    def show_img(self,arg):
        self.image_sig.emit(arg)

    def show_aoi_img(self,arg):
        self.aoi_image_sig.emit(arg)

    def ThreeD_image(self,arg):
        self._3d_image_sig.emit(arg)

    def aoi_tablewidget(self,arg):
        self.aoi_row_sig.emit(arg)

    def ai_tablewidget(self,arg):
        self.ai_row_sig.emit(arg)

    def clear_widget(self):
        self.clear_sig.emit()

    def roi_groupbox(self,arg):
        self.aoi_box_sig.emit(arg)

    def Stop(self):
        try:
            self.timer.stop()
            self.work.quit()
            self.work.wait()
            del self.work
        except Exception as e:
            print(e)


if __name__=='__main__':
    th = Win_FileWatcher()
    th.add_watchlist('','')